package com.etl.dataflow.executor.core;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.collection.ListUtil;
import com.etl.dataflow.common.entity.JdbcProperties;
import com.etl.dataflow.common.exception.EtlException;
import com.etl.dataflow.executor.element.Column;
import com.etl.dataflow.executor.element.Row;
import com.etl.dataflow.executor.exception.UnsupportedFieldTypeException;
import com.etl.dataflow.executor.util.JdbcQuerySaveUtil;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Types;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author dx
 * @since 2022/4/14
 */
@Slf4j
public class OracleWriter implements RdbmsWriter {

    /**
     * 速度测量时间间隔(s)
     */
    private static final int SPEED_MEASURE_INTERVAL = 5;
    protected final String tableName;
    protected final int batchSize = 1000;
    protected final ChannelReader channelReader;
    protected final JdbcProperties jdbcProperties;
    protected boolean shutdown = false;

    protected final TaskContext taskContext;
    private String sql4saveRow = null;
    private OnWriteFinishedListener onWriteFinishedListener;

    private OnFatalErrorListener onFatalErrorListener;

    private AbstractAfterWriteInterceptor afterWriteInterceptor;

    public OracleWriter(String tableName, JdbcProperties jdbcProperties, ChannelReader channelReader) {
        this.tableName = tableName;
        this.channelReader = channelReader;
        this.jdbcProperties = jdbcProperties;

        Optional<TaskContext> taskContextOptional = TaskContextHolder.getTaskContext();
        assert taskContextOptional.isPresent();
        taskContext = taskContextOptional.get();
    }

    @Override
    public void startWrite() {
        Connection connection = null;
        try {
            connection = JdbcQuerySaveUtil.getConnection(jdbcProperties);
            final List<Row> buffer = new ArrayList<>();
            boolean isFirstRow = true;
            long lastMeasureTime = System.currentTimeMillis();
            double interval;
            long rowNumber = 0L;
            long byteSize = 0L;
            while (true) {
                Row row = channelReader.readOneRow();
                // 读取不到记录说明已经全部结束
                if (row == null || shutdown) {
                    break;
                }
                // 第一次插入需要生成sql
                if (isFirstRow) {
                    generateInsertSql(row);
                    isFirstRow = false;
                }
                buffer.add(row);
                if (buffer.size() >= batchSize) {
                    saveBatch(connection, buffer);
                    buffer.clear();
                }

                rowNumber++;
                byteSize += row.getContentByteSize();

                interval = (System.currentTimeMillis() - lastMeasureTime) / 1e+3;
                if (interval >= SPEED_MEASURE_INTERVAL) {
                    taskContext.getTaskStatisticInfo().getReadBytesPerSecond().set((long) (byteSize / interval));
                    byteSize = 0;
                    taskContext.getTaskStatisticInfo().getReadRowsPerSecond().set((long) (rowNumber / interval));
                    rowNumber = 0;
                    lastMeasureTime = System.currentTimeMillis();
                }
            }
            if (shutdown) {
                if (this.onWriteFinishedListener != null) {
                    this.onWriteFinishedListener.onFinished();
                }
                return;
            }
            if (CollUtil.isNotEmpty(buffer)) {
                saveBatch(connection, buffer);
            }
            log.debug("写入结束");
            invokeAfterWriteInterceptor(null, null);
            if (this.onWriteFinishedListener != null) {
                this.onWriteFinishedListener.onFinished();
            }
        } catch (Exception e) {
            if (onFatalErrorListener != null) {
                onFatalErrorListener.onError(e);
            }
            throw new EtlException(e);
        } finally {
            if (connection != null) {
                JdbcQuerySaveUtil.closeDbResources(null, connection);
            }
            TaskContextHolder.removeTaskContext();
        }
    }

    @Override
    public void setOnWriteFinishedListener(OnWriteFinishedListener listener) {
        this.onWriteFinishedListener = listener;
    }

    @Override
    public void setOnFatalErrorListener(OnFatalErrorListener listener) {
        this.onFatalErrorListener = listener;
    }

    @Override
    public void setAfterWriteInterceptor(AbstractAfterWriteInterceptor interceptor) {
        this.afterWriteInterceptor = interceptor;
    }

    @Override
    public void exit() {
        this.shutdown = true;
    }

    private void generateInsertSql(Row row) {
        if (sql4saveRow != null) {
            return;
        }
        List<String> fields = new ArrayList<>();
        for (int i = 1; i <= row.getColumnNumber(); i++) {
            fields.add(row.getColumn(i).getName());
        }
        this.sql4saveRow = String.format(
                "insert /*+Append*/ into %s(%s) values(%s)",
                this.tableName,
                CollUtil.join(fields, ","),
                generatePlaceHolder(fields.size())
        );
    }

    private String generatePlaceHolder(int count) {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < count; i++) {
            sb.append("?,");
        }
        sb.deleteCharAt(sb.length() - 1);
        return sb.toString();
    }

    @SneakyThrows
    protected void saveBatch(Connection connection, List<Row> rows) {
        PreparedStatement preparedStatement = null;
        AtomicInteger failedCount = new AtomicInteger(0);
        try {
            connection.setAutoCommit(false);
            preparedStatement = connection.prepareStatement(this.sql4saveRow);
            for (Row row : rows) {
                fillPreparedStatement(preparedStatement, row);
                preparedStatement.addBatch();
            }
            preparedStatement.executeBatch();
            connection.commit();
            // 执行写后拦截器
            invokeAfterWriteInterceptor(rows, ListUtil.empty());
        } catch (SQLException e) {
            System.out.println("回滚此次写入, 采用每次写入一行方式提交. 因为:" + e.getMessage());
            connection.rollback();
            int failedSize = doOneInsert(connection, rows);
            if (failedSize > 0) {
                failedCount.getAndAdd(failedSize);
            }
        } catch (Exception e) {
            throw new EtlException(e);
        } finally {
            JdbcQuerySaveUtil.closeDbResources(preparedStatement, null);

            TaskContextHolder.getTaskContext().ifPresent(taskContext -> {
                // 插入成功的数量 = 总数量 - 插入失败数量
                taskContext.getTaskStatisticInfo()
                        .getWrittenRowNumber().addAndGet(rows.size() - failedCount.get());
                // 已处理条数
                taskContext.getTaskStatisticInfo()
                        .getProcessedRowNumber().addAndGet(rows.size());
            });
        }
    }

    protected int doOneInsert(Connection connection, List<Row> rows) {
        List<Row> writtenRows = new ArrayList<>();
        List<Row> failedRows = new ArrayList<>();
        PreparedStatement preparedStatement = null;
        try {
            connection.setAutoCommit(true);
            preparedStatement = connection.prepareStatement(this.sql4saveRow);
            for (Row row : rows) {
                try {
                    fillPreparedStatement(preparedStatement, row);
                    preparedStatement.execute();
                    writtenRows.add(row);
                } catch (SQLException e) {
                    // 记录未插入的脏数据交由拦截器自行实现
                    failedRows.add(row);
                    log.error("ERROR MESSAGE: {} ERROR SQL ===> {}, params ===> {}", e.getMessage(), this.sql4saveRow, row);
                } finally {
                    // 最后不要忘了关闭 preparedStatement
                    preparedStatement.clearParameters();
                }
            }
        } catch (Exception e) {
            throw new EtlException(e);
        } finally {
            JdbcQuerySaveUtil.closeDbResources(preparedStatement, null);
        }
        int failedSize = failedRows.size();
        // 执行写后拦截器
        invokeAfterWriteInterceptor(writtenRows, failedRows);
        return failedSize;
    }

    private void fillPreparedStatement(PreparedStatement preparedStatement, Row row)
            throws SQLException {
        for (int i = 1; i <= row.getColumnNumber(); i++) {
            fillPreparedStatementColumnType(preparedStatement, i, row.getColumn(i));
        }
    }


    private void fillPreparedStatementColumnType(PreparedStatement preparedStatement, int columnIndex,
                                                 Column<?> column) throws SQLException {
        switch (column.getJdbcType()) {
            case Types.CHAR:
            case Types.NCHAR:
            case Types.CLOB:
            case Types.NCLOB:
            case Types.VARCHAR:
            case Types.LONGVARCHAR:
            case Types.NVARCHAR:
            case Types.LONGNVARCHAR:

            case Types.BOOLEAN:
                preparedStatement.setString(columnIndex, column.asString());
                break;

            case Types.SMALLINT:
            case Types.INTEGER:
            case Types.BIGINT:
            case Types.NUMERIC:
            case Types.DECIMAL:
            case Types.FLOAT:
            case Types.REAL:
            case Types.TINYINT:
            case Types.DOUBLE: {
                preparedStatement.setString(columnIndex, column.asString());
                break;
            }

            case Types.DATE: {
                preparedStatement.setDate(columnIndex, column.asDate());
                break;
            }

            case Types.TIME: {
                java.sql.Time sqlTime = new java.sql.Time(column.asDate().getTime());
                preparedStatement.setTime(columnIndex, sqlTime);
                break;
            }

            case Types.TIMESTAMP: {
                preparedStatement.setTimestamp(columnIndex, column.asTimestamp());
                break;
            }

            case Types.BINARY:
            case Types.VARBINARY:
            case Types.BLOB:
            case Types.LONGVARBINARY:
                preparedStatement.setBytes(columnIndex, column.asBytes());
                break;
            default:
                throw new UnsupportedFieldTypeException("不支持的插入数据类型：" + column.asString());
        }
    }

    protected void invokeAfterWriteInterceptor(List<Row> writtenRows, List<Row> failedRows) {
        if (this.afterWriteInterceptor == null) {
            return;
        }
        if (writtenRows == null && failedRows == null) {
            afterWriteInterceptor.onWriteFinished();
            return;
        }
        afterWriteInterceptor.onAfterWrite(writtenRows, failedRows);
    }
}
